# get_data.py
# -*- coding: utf-8 -*-
"""
数据接口模块，提供统一的数据获取接口
适配通达信数据源
"""

import pandas as pd
from datetime import datetime, timedelta
import numpy as np

# 导入通达信数据库接口
try:
    from pytdx.hq import TdxHq_API
    from pytdx.params import TDXParams
    TDX_AVAILABLE = True
except ImportError:
    print("⚠️  pytdx 未安装，请运行: pip install pytdx")
    TDX_AVAILABLE = False

# 通达信服务器列表（更多服务器地址）
TDX_SERVERS = [
    ('180.153.18.170', 7709), # 备用服务器8
    ('202.108.253.130', 7709), # 备用服务器9
]

def get_tdx_api():
    """获取通达信API连接（带重试机制）"""
    if not TDX_AVAILABLE:
        return None
    
    import time
    
    api = TdxHq_API()
    # 尝试连接服务器
    for host, port in TDX_SERVERS:
        try:
            print(f"尝试连接通达信服务器: {host}:{port}")
            if api.connect(host, port, time_out=10):
                print(f"✅ 成功连接到 {host}:{port}")
                return api
            else:
                print(f"❌ 连接失败: {host}:{port}")
        except Exception as e:
            print(f"❌ 连接异常 {host}:{port}: {e}")
            continue
        
        # 每次尝试后稍微等待
        time.sleep(0.5)
    
    print("❌ 所有通达信服务器连接失败")
    return None

def convert_stock_code(symbol):
    """转换股票代码格式"""
    if symbol.endswith('.XSHG'):
        # 上海交易所
        stock_code = symbol.replace('.XSHG', '')
        market = 1  # 上海市场
    elif symbol.endswith('.XSHE'):
        # 深圳交易所
        stock_code = symbol.replace('.XSHE', '')
        market = 0  # 深圳市场
    else:
        # 根据代码判断市场
        stock_code = symbol
        if symbol.startswith('6'):
            market = 1  # 上海
        else:
            market = 0  # 深圳
    
    return stock_code, market

def get_price(symbol, start_date=None, end_date=None, count=None, frequency='daily', fields=['close']):
    """
    获取股票价格数据（使用通达信数据源）
    
    参数:
    symbol: 股票代码，如 '159915.XSHE'
    start_date: 开始日期，格式 'YYYYMMDD'
    end_date: 结束日期，格式 'YYYYMMDD'
    count: 获取数据的数量
    frequency: 频率，如 'daily'
    fields: 需要获取的字段列表
    
    返回:
    DataFrame: 包含指定字段的价格数据
    """
    if not TDX_AVAILABLE:
        print(f"通达信库未安装，无法获取数据 {symbol}")
        return None
        
    try:
        # 转换股票代码格式
        stock_code, market = convert_stock_code(symbol)
        
        # 获取通达信API连接
        api = get_tdx_api()
        if api is None:
            print(f"无法连接通达信服务器 {symbol}")
            return None
        
        # 处理日期和数量参数
        if count is None:
            count = 800  # 默认获取800天数据
        
        # 通达信按天获取数据，每次最多800条
        all_data = []
        remaining_count = count
        start_pos = 0
        
        while remaining_count > 0:
            batch_count = min(remaining_count, 800)+1
            
            # 获取K线数据
            data = api.get_security_bars(
                category=TDXParams.KLINE_TYPE_DAILY,  # 日K线
                market=market,
                code=stock_code,
                start=start_pos,
                count=batch_count
            )
            
            if not data:
                break
                
            all_data.extend(data)
            remaining_count -= len(data)
            start_pos += len(data)
            
            # 如果返回的数据少于请求的数量，说明没有更多数据了
            if len(data) < batch_count:
                break
        
        # 关闭连接
        api.disconnect()
        
        if not all_data:
            print(f"未获取到数据 {symbol}")
            return None
        
        # 转换为DataFrame
        if not all_data:
            print(f"未获取到数据 {symbol}")
            return None
            
        try:
            if all_data:
                # 确保数据是列表格式
                if isinstance(all_data[0], dict):
                    df = pd.DataFrame(all_data)
                else:
                    # 如果是其他格式，尝试转换
                    df = pd.DataFrame(all_data)
            else:
                print(f"数据为空 {symbol}")
                return None
        except Exception as e:
            print(f"DataFrame创建失败 {symbol}: {e}")
            print(f"数据样本: {all_data[0] if all_data else 'None'}")
            return None
        
        if df.empty:
            print(f"数据为空 {symbol}")
            return None
        
        # 重命名列
        column_mapping = {
            'datetime': 'date',
            'open': 'open',
            'close': 'close',
            'high': 'high',
            'low': 'low',
            'vol': 'volume',
            'amount': 'amount'
        }
        
        # 只重命名存在的列
        existing_columns = {k: v for k, v in column_mapping.items() if k in df.columns}
        if existing_columns:
            df = df.rename(columns=existing_columns)
        else:
            print(f"未找到预期的列名 {symbol}, 实际列名: {list(df.columns)}")
            return None
        
        # 处理日期 - 通达信返回标准日期格式
        if 'date' in df.columns:
            # 直接转换，通达信返回的是标准格式
            df['date'] = pd.to_datetime(df['date'], errors='coerce')
            # 删除无效日期的行并设置索引
            df = df.dropna(subset=['date'])
            df = df.set_index('date')
            df = df.sort_index()  # 按日期排序
        
        # 优化日期过滤逻辑：优先满足count需求
        if count:
            # 如果指定了count，先按count取数据
            if len(df) > count:
                df = df.tail(count)
            
            # 然后再按end_date过滤（如果指定了的话）
            if end_date:
                if isinstance(end_date, str):
                    end_date = pd.to_datetime(end_date)
                elif isinstance(end_date, datetime):
                    end_date = pd.to_datetime(end_date)
                
                # 只保留end_date及之前的数据
                df = df[df.index <= end_date]
                
                # 如果过滤后数据不足count，尝试获取更多数据
                if len(df) < count:
                    print(f"  📊 日期过滤后数据不足，重新获取更多数据...")
                    # 重新获取更多数据
                    api2 = get_tdx_api()
                    if api2:
                        try:
                            # 获取更多数据以满足count需求
                            more_data = api2.get_security_bars(
                                category=TDXParams.KLINE_TYPE_DAILY,
                                market=market,
                                code=stock_code,
                                start=0,
                                count=min(count * 2, 800)  # 获取count的2倍数据
                            )
                            api2.disconnect()
                            
                            if more_data:
                                # 重新处理数据
                                df_more = pd.DataFrame(more_data)
                                # 只重命名存在的列
                                existing_columns = {k: v for k, v in column_mapping.items() if k in df_more.columns}
                                if existing_columns:
                                    df_more = df_more.rename(columns=existing_columns)
                                if 'date' in df_more.columns:
                                    df_more['date'] = pd.to_datetime(df_more['date'], errors='coerce')
                                    df_more = df_more.dropna(subset=['date']).set_index('date')
                                    df_more = df_more.sort_index()
                                    
                                    # 过滤到end_date并取count条数据
                                    df_filtered = df_more[df_more.index <= end_date]
                                    if len(df_filtered) >= count:
                                        df = df_filtered.tail(count)
                        except Exception as e:
                            print(f"获取更多数据时出错: {e}")
                            pass
        else:
            # 如果没有指定count，按日期范围过滤
            if start_date or end_date:
                if start_date:
                    if isinstance(start_date, str):
                        start_date = pd.to_datetime(start_date)
                    elif isinstance(start_date, datetime):
                        start_date = pd.to_datetime(start_date)
                    df = df[df.index >= start_date]
                
                if end_date:
                    if isinstance(end_date, str):
                        end_date = pd.to_datetime(end_date)
                    elif isinstance(end_date, datetime):
                        end_date = pd.to_datetime(end_date)
                    df = df[df.index <= end_date]
        
        # 只返回需要的字段
        if fields:
            available_fields = [f for f in fields if f in df.columns]
            if available_fields:
                df = df[available_fields]
            else:
                print(f"没有找到请求的字段 {fields}")
        
        return df if len(df) > 0 else None
        
    except Exception as e:
        print(f"获取股票数据失败 {symbol}: {e}")
        import traceback
        traceback.print_exc()
        return None

def get_index_price(symbol, start_date=None, end_date=None):
    """
    获取指数价格数据（使用通达信数据源）
    
    参数:
    symbol: 指数代码，如 '000300.XSHG'
    start_date: 开始日期，格式 'YYYYMMDD'
    end_date: 结束日期，格式 'YYYYMMDD'
    
    返回:
    DataFrame: 包含指数价格数据
    """
    if not TDX_AVAILABLE:
        print(f"通达信库未安装，无法获取指数数据 {symbol}")
        return None
        
    try:
        # 处理指数代码格式
        if symbol == '000300.XSHG':
            index_code = '000300'
            market = 1  # 上海市场
        elif symbol == '000001.XSHG':
            index_code = '000001'
            market = 1  # 上证指数
        elif symbol == '399001.XSHE':
            index_code = '399001'
            market = 0  # 深证成指
        else:
            # 默认处理方式
            if symbol.endswith('.XSHG'):
                index_code = symbol.replace('.XSHG', '')
                market = 1
            elif symbol.endswith('.XSHE'):
                index_code = symbol.replace('.XSHE', '')
                market = 0
            else:
                index_code = symbol
                market = 1  # 默认上海
        
        
        # 获取通达信API连接
        api = get_tdx_api()
        if api is None:
            print(f"无法连接通达信服务器 {symbol}")
            return None
        
        # 获取指数K线数据，默认获取800天
        count = 800
        if start_date or end_date:
            # 如果指定了日期范围，计算需要的天数
            if end_date:
                if isinstance(end_date, str):
                    end_date = pd.to_datetime(end_date)
                elif isinstance(end_date, datetime):
                    end_date = pd.to_datetime(end_date)
            if start_date:
                if isinstance(start_date, str):
                    start_date = pd.to_datetime(start_date)
                elif isinstance(start_date, datetime):
                    start_date = pd.to_datetime(start_date)
                # 计算需要的天数
                if end_date:
                    count = (end_date - start_date).days + 50  # 多取一些缓冲
                else:
                    count = 800
            else:
                count = 800
        
        all_data = []
        remaining_count = count
        start_pos = 0
        
        while remaining_count > 0:
            batch_count = min(remaining_count, 800)
            
            # 获取指数K线数据
            data = api.get_index_bars(
                category=TDXParams.KLINE_TYPE_DAILY,  # 日K线
                market=market,
                code=index_code,
                start=start_pos,
                count=batch_count
            )
            
            if not data:
                break
                
            all_data.extend(data)
            remaining_count -= len(data)
            start_pos += len(data)
            
            # 如果返回的数据少于请求的数量，说明没有更多数据了
            if len(data) < batch_count:
                break
        
        # 关闭连接
        api.disconnect()
        
        if not all_data:
            print(f"未获取到指数数据 {symbol}")
            return None
        
        
        # 转换为DataFrame
        if not all_data:
            print(f"未获取到指数数据 {symbol}")
            return None
            
        try:
            if all_data:
                # 确保数据是列表格式
                if isinstance(all_data[0], dict):
                    df = pd.DataFrame(all_data)
                else:
                    # 如果是其他格式，尝试转换
                    df = pd.DataFrame(all_data)
            else:
                print(f"指数数据为空 {symbol}")
                return None
        except Exception as e:
            print(f"指数DataFrame创建失败 {symbol}: {e}")
            print(f"指数数据样本: {all_data[0] if all_data else 'None'}")
            return None
        
        if df.empty:
            print(f"指数数据为空 {symbol}")
            return None
        
        # 重命名列
        column_mapping = {
            'datetime': 'date',
            'open': 'open',
            'close': 'close',
            'high': 'high',
            'low': 'low',
            'vol': 'volume',
            'amount': 'amount'
        }
        
        # 只重命名存在的列
        existing_columns = {k: v for k, v in column_mapping.items() if k in df.columns}
        if existing_columns:
            df = df.rename(columns=existing_columns)
        else:
            print(f"指数未找到预期的列名 {symbol}, 实际列名: {list(df.columns)}")
            return None
        
        # 处理日期 - 通达信返回标准日期格式
        if 'date' in df.columns:
            # 直接转换，通达信返回的是标准格式 "2025-09-22 15:00"
            df['date'] = pd.to_datetime(df['date'], errors='coerce')
            # 删除无效日期的行并设置索引
            df = df.dropna(subset=['date']).set_index('date')
            df = df.sort_index()
            
            # 应用日期过滤
            if start_date or end_date:
                if start_date:
                    df = df[df.index >= start_date]
                if end_date:
                    df = df[df.index <= end_date]
        else:
            print(f"指数未找到date列 {symbol}")
            return None
        
        return df if len(df) > 0 else None
        
    except Exception as e:
        print(f"指数数据处理失败 {symbol}: {e}")
        return None

def get_security_info(symbol):
    """
    获取证券信息
    
    参数:
    symbol: 证券代码
    
    返回:
    SecurityInfo对象
    """
    class SecurityInfo:
        def __init__(self, code, display_name):
            self.code = code
            self.display_name = display_name
    
    # 代码到名称的映射
    name_mapping = {
        '159915.XSHE': '创业板',
        '159915': '创业板',
        '510300.XSHG': '沪深300',
        '510300': '沪深300',
        '510500.XSHG': '中证500',
        '510500': '中证500',
        '000300.XSHG': '沪深300',
        '000300': '沪深300',
        '000001.XSHG': '上证指数',
        '000001': '上证指数',
        '399001.XSHE': '深证成指',
        '399001': '深证成指',
    }
    
    try:
        # 获取股票名称
        name = name_mapping.get(symbol, symbol)
        
        # 如果使用通达信API，可以尝试获取更详细信息
        if TDX_AVAILABLE:
            try:
                api = get_tdx_api()
                if api:
                    stock_code, market = convert_stock_code(symbol)
                    # 获取股票信息
                    info = api.get_security_info(market, stock_code)
                    api.disconnect()
                    if info and len(info) > 0:
                        name = info[0].get('name', name)
            except:
                pass  # 如果获取失败，使用默认名称
        
        return SecurityInfo(symbol, name)
        
    except Exception as e:
        print(f"获取股票信息失败 {symbol}: {e}")
        return SecurityInfo(symbol, symbol)

def get_trade_days(end_date=None, count=5, start_date=None):
    """
    获取交易日列表（简化版本，排除周末）
    
    参数:
    end_date: 结束日期
    count: 交易日数量
    start_date: 开始日期
    
    返回:
    list: 交易日列表
    """
    try:
        if not end_date:
            end_date = datetime.now()
        elif isinstance(end_date, str):
            end_date = datetime.strptime(end_date, '%Y%m%d')
            
        # 简单的交易日生成：排除周末
        trade_days = []
        current_date = end_date
        days_back = 0
        
        while len(trade_days) < count and days_back < count * 3:  # 最多回溯count*3天
            # 排除周末（周六=5，周日=6）
            if current_date.weekday() < 5:
                trade_days.append(current_date)
            current_date = current_date - timedelta(days=1)
            days_back += 1
        
        # 按时间顺序返回
        trade_days.reverse()
        return trade_days[-count:] if len(trade_days) >= count else trade_days
        
    except Exception as e:
        print(f"获取交易日失败: {e}")
        # 出错时返回简单日期列表
        if not end_date:
            end_date = datetime.now()
        return [end_date - timedelta(days=x) for x in range(count-1, -1, -1)]
